import numpy as np
import scipy.special as sc_special

# 本次修改将参数前置
nest_num = 50
nest_dim = 2
iter_num = 10000
pa = 0.25
beta = 1.5  # 莱维飞行公式的β
step_coefficient = 0.01  # 步长控制因子

# 本次修改在遗弃函数中加入差分策略,称为双策略遗弃，交叉选择最优的候选解
version = '1.3.0'


# 生成鸟窝
def generate_neste(nest_num, nest_dim, lower_boundary, upper_boundary):
    lower_boundary = np.array(lower_boundary)
    upper_boundary = np.array(upper_boundary)

    nests = np.empty((nest_num, nest_dim))

    # 遍历生成鸟窝，此问题相当于生成x，y
    for each_nest in range(nest_num):
        # 下边界 + 系数 * (上-下)
        nests[each_nest] = lower_boundary + np.array([np.random.rand() for _ in range(nest_dim)]) * (
                upper_boundary - lower_boundary)

    return nests


# 适应度计算函数
def calc_fitness(fit_func, nests):
    finess = np.empty(nest_num)

    # 遍历nests 计算 适应度 填入 finess
    for each_nest in range(nest_num):
        finess[each_nest] = fit_func(nests[each_nest])

    return finess


# 莱维飞行函数,完全按照论文复现
def levy_flight(nest_num, nest_dim, beta):
    sigma_u = (sc_special.gamma(1 + beta) * np.sin(np.pi * beta / 2) / (
            sc_special.gamma((1 + beta) / 2) * beta * (2 ** ((beta - 1) / 2)))) ** (1 / beta)
    sigma_v = 1

    # 正态分布参数：均值，标准差，shape
    u = np.random.normal(0, sigma_u, (nest_num, nest_dim))
    v = np.random.normal(0, sigma_v, (nest_num, nest_dim))

    steps = u / ((np.abs(v)) ** (1 / beta))

    return steps


# 更新鸟窝函数
def update_nests(fit_func, lower_boundary, upper_boundary, nests, best_nest, fitness, step_coefficient):
    lower_boundary = np.array(lower_boundary)
    upper_boundary = np.array(upper_boundary)

    steps = levy_flight(nest_num, nest_dim, beta)

    new_nests = nests.copy()

    for each_nest in range(nest_num):
        # 步长的计算，步长系数*步长*(鸟窝-最佳鸟窝)
        step_size = step_coefficient * steps[each_nest] * (nests[each_nest] - best_nest)

        # 随机生成方向
        # step_direction = np.random.rand(nest_dim)
        # 更新nest,原始步长+步长*方向
        # new_nests[each_nest] += step_size * step_direction

        # 没有随机方向的
        new_nests[each_nest] += step_size
        # 新生成的nest小于下界即取下界
        new_nests[each_nest][new_nests[each_nest] < lower_boundary] = lower_boundary[
            new_nests[each_nest] < lower_boundary]
        # 新生成的nest大于上界即取上界
        new_nests[each_nest][new_nests[each_nest] > upper_boundary] = upper_boundary[
            new_nests[each_nest] > upper_boundary]

    # 求新的nest的适应值
    new_fitness = calc_fitness(fit_func, new_nests)
    # 新的适应值大于 旧适应值 取新适应值 对应的nest
    nests[new_fitness > fitness] = new_nests[new_fitness > fitness]

    return nests


# 遗弃鸟窝的函数，也就是二次更新函数，没有加论文中的跃进函数
def abandon_nests(nests, lower_boundary, upper_boundary, pa):
    lower_boundary = np.array(lower_boundary)
    upper_boundary = np.array(upper_boundary)

    for each_nest in range(nest_num):
        # 生成随机数来控制遗弃的概率,生成随机数的范围为0-1
        if (np.random.rand() < pa):
            # 步长 = 随机数* (俩个随机nest的差)，除了加跃进函数，完全严格按照论文复现
            step_size = np.random.rand() * (
                    nests[np.random.randint(0, nest_num)] - nests[np.random.randint(0, nest_num)])
            nests[each_nest] += step_size
            # 越界取值判断，同理，超过上下界取上下界
            nests[each_nest][nests[each_nest] < lower_boundary] = lower_boundary[nests[each_nest] < lower_boundary]
            nests[each_nest][nests[each_nest] > upper_boundary] = upper_boundary[nests[each_nest] > upper_boundary]

        return nests


# 俩种策略的 鸟窝遗弃，也就是二次更新，严格按照论文复现
def bsrw(nests, pa, lower_boundary, upper_boundary):
    lower_boundary = np.array(lower_boundary)
    upper_boundary = np.array(upper_boundary)
    for each_nest in range(nest_num):
        a = nest_num - each_nest + 1
        # 定义交叉因子中的C
        C = 2.718281828459045 ** (1 - nest_num / a)

        if (np.random.rand() < pa):
            # 生成三个nests的随机位置
            index1 = np.random.randint(0, nest_num)
            index2 = np.random.randint(0, nest_num)
            index3 = np.random.randint(0, nest_num)

            # 局部随机(原始)
            step_size1 = np.random.rand() * (nests[index2] - nests[index3])
            # 随机游走产生的候选解
            v1 = nests[each_nest] + step_size1

            # DE算法，F变异因子取0.6
            step_size2 = 0.6 * (nests[index1] - nests[index2])
            # 差分算法产生的候选解
            v2 = nests[each_nest] + step_size2

            # 交叉选择策略
            if ((np.random.rand() < (0.15 * 4 ** C)) or (each_nest == np.random.randint(0, 2))):
                nests[each_nest] = v1
            else:
                nests[each_nest] = v2
            # 越界取值判断，同理，超过上下界取上下界
            nests[each_nest][nests[each_nest] < lower_boundary] = lower_boundary[nests[each_nest] < lower_boundary]
            nests[each_nest][nests[each_nest] > upper_boundary] = upper_boundary[nests[each_nest] > upper_boundary]
    return nests


def run(nest_num, nest_dim, fit_func, lower_boundary, upper_boundary, iter_num, pa, beta,
        step_coefficient):
    # 生成鸟窝nest-x,y
    nests = generate_neste(nest_num, nest_dim, lower_boundary, upper_boundary)

    # 计算适应度
    fitness = calc_fitness(fit_func, nests)

    # 记录最佳适应度的位置
    best_nest_index = np.argmin(fitness)

    # 记录最佳适应度的值
    best_fitness = fitness[best_nest_index]
    # 记录最佳适应度对应的nest的值(鸟窝)
    best_nest = nests[best_nest_index].copy()

    # 初始化一个最佳适应度，相当于上一次的最佳适应度
    min_fitness = 0

    # 我理解的s实际上是一种记录量，用来记录莱维飞行生成的nest的最佳适应度值是否发生了变化，如果最佳适应度与上次一样就+1
    s = 0
    # M是控制变量的临界条件，当发生7次最佳适应度值无变化时我们就采取差分进化的思想，差分进化完成后将s归0
    M = 7
    for _ in range(iter_num):
        # 如果当前最佳适应度等于之前的最佳适应度，s+1
        if (min_fitness == best_fitness):
            s += 1
            continue
        if s <= M:
            nests = update_nests(fit_func, lower_boundary, upper_boundary, nests, best_nest, fitness, step_coefficient)
        else:
            # 差分进化算法的论文复现 变异-交叉-选择
            # 随机取变异位置
            index = np.random.randint(0, nest_num)
            # 生成变异结果也就是候选解
            v = nests[index] + 0.6 * (
                    nests[np.random.randint(0, nest_num)] - nests[np.random.randint(0, nest_num)])
            # 进行变异结果和原鸟窝的二项式交叉,交叉因子取0.1
            if ((np.random.rand(1) <= 0.1) or (index == np.random.randint(0, 2))):
                u = v
            else:
                u = nests[index]
            # 论文给的是最小优化问题，
            # 选择策略

            u_fit = calc_fitness(fit_func, u)
            x_fit = calc_fitness(fit_func, nests[index])
            if (u_fit <= x_fit):
                nests[index] = u
            else:
                pass
            # s置为0
            s = 0

        # 混合差分来进行鸟窝遗弃
        nests = bsrw(nests, pa, lower_boundary, upper_boundary)
        # 计算适应度
        fitness = calc_fitness(fit_func, nests)

        # 记录最小适应度所在位置
        # max_nest_index = np.argmax(fitness)
        min_nest_index = np.argmin(fitness)

        # 最小适应度的值
        min_fitness = fitness[min_nest_index]
        # 最小适应度对应的nest -- x，y
        min_nest = nests[min_nest_index]

        # 判断，如果该次的小于适应度大于之前的，就换掉
        if (min_fitness < best_fitness):
            best_nest = min_nest.copy()
            best_fitness = min_fitness
        # 返回最佳nest--xy，和最小适应度
    return (best_nest, best_fitness)


if __name__ == '__main__':
    def fit_func(nest):
        x, y = nest
        return 200 - (x ** 2 + y - 11) ** 2 - (x + y ** 2 - 7) ** 2


    best_nest, best_fitness = run(nest_num, nest_dim, fit_func, [-3, -3], [3, 3], iter_num, pa, beta, step_coefficient)

    print('最小值为:%.5f, 在(%.5f, %.5f)处取到!' % (best_fitness, best_nest[0], best_nest[1]))
